#include "ModelTorch.h"

namespace Bublik
{

 ModelTorch::ModelTorch(std::string& FileName)
 {
 const bool cuda_is_available = torch::cuda::is_available();
 _device = cuda_is_available ? torch::kCUDA : torch::kCPU;
 _imgSize_x = 640;
 _imgSize_y = 640;
 _model = torch::jit::load(FileName, _device);
 }

 cv::Mat ModelTorch::getDetection(std::string& ImageFileName, float Inputconf)
 {
 _inputImage = cv::imread(ImageFileName, cv::ImreadModes::IMREAD_COLOR);
 cv::cvtColor(_inputImage, _inputImage, cv::COLOR_BGR2RGB);
 cv::Mat image_transfomed;

 _inputImage = PrepImg(_inputImage, _imgSize_x, _imgSize_y);
 _inputImage.copyTo(image_transfomed);

 torch::Tensor tensor_image = torch::from_blob(image_transfomed.data,
 { image_transfomed.rows, image_transfomed.cols,3 }, torch::kByte);
 tensor_image = tensor_image.permute({ 2,0,1 });
 tensor_image = tensor_image.toType(torch::kFloat);
 tensor_image = tensor_image.div(255);
 tensor_image = tensor_image.unsqueeze(0);
 tensor_image = tensor_image.to(_device);
 auto output = _model.forward({ tensor_image }).toTuple();

 at::Tensor scores = output->elements()[0].toTensor();

 std::vector<std::vector<int>> boxes;
 float confThres = Inputconf;
 for (unsigned short ibatch = 0; ibatch < scores.sizes()[0]; ibatch++) {
 for (unsigned short ibox = 0; ibox < scores.sizes()[2]; ibox++) {
 float conf = scores[ibatch][4][ibox].item<float>();
 if (conf >= confThres) {
 unsigned short
 cx = scores[ibatch][0][ibox].item<int>(),
 cy = scores[ibatch][1][ibox].item<int>(),
 w = scores[ibatch][2][ibox].item<int>(),
 h = scores[ibatch][3][ibox].item<int>();
 unsigned short
 x1 = cx - w / 2,
 y1 = cy - h / 2,
 x2 = cx + w / 2,
 y2 = cy + h / 2;

 boxes.push_back({static_cast<int>(conf*100), x1, y1, x2, y2});

 }
 }
 }
 std::sort(boxes.begin(), boxes.end(), [](const std::vector<int>& b1, const std::vector<int>& b2) {return b1[0] > b2[0]; });
 unsigned short
 x1 = static_cast<unsigned short>(boxes[0][1]),
 y1 = static_cast<unsigned short>(boxes[0][2]),
 x2 = static_cast<unsigned short>(boxes[0][3]),
 y2 = static_cast<unsigned short>(boxes[0][4]);
 cv::Scalar rectColor(0, 192, 0);
 cv::rectangle(_inputImage, { x1,y1 }, { x2,y2 }, rectColor, 2);
 return _inputImage;
 }

 cv::Mat ModelTorch::getSegmentation(std::string& ImageFileName, float Inputconf)
 {
 _inputImage = cv::imread(ImageFileName, cv::ImreadModes::IMREAD_COLOR);
 cv::cvtColor(_inputImage, _inputImage, cv::COLOR_BGR2RGB);
 cv::Mat image_transfomed;

 _inputImage = PrepImg(_inputImage, _imgSize_x, _imgSize_y);
 _inputImage.copyTo(image_transfomed);

 torch::Tensor tensor_image = torch::from_blob(image_transfomed.data,
 { image_transfomed.rows, image_transfomed.cols,3 }, torch::kByte);
 tensor_image = tensor_image.permute({ 2,0,1 });
 tensor_image = tensor_image.toType(torch::kFloat);
 tensor_image = tensor_image.div(255);
 tensor_image = tensor_image.unsqueeze(0);
 tensor_image = tensor_image.to(_device);
 auto output = _model.forward({ tensor_image }).toTuple();

 at::Tensor scores = output->elements()[0].toTensor();
 at::Tensor descriptors = output->elements()[1].toTensor();

 std::vector<std::vector<int>> segmentation;
 cv::Mat img160(160, 160, CV_32FC1);
 cv::Mat img160u;
 float confThres = Inputconf;
 for (unsigned short ibatch = 0; ibatch < scores.sizes()[0]; ibatch++) {
 for (unsigned short ibox = 0; ibox < scores.sizes()[2]; ibox++) {
 float conf = scores[ibatch][4][ibox].item<float>();
 if (conf >= confThres) {
 int k_max = 0;
 float v_max = scores[ibatch][k_max + 5][ibox].item<float>();
 for (int k = 0; k < 32; k++)
 {
 float v = scores[ibatch][k + 5][ibox].item<float>();
 if (v > v_max)
 {
 v_max = v;
 k_max = k;
 }
 }
 unsigned short
 cx = scores[ibatch][0][ibox].item<int>(),
 cy = scores[ibatch][1][ibox].item<int>(),
 w = scores[ibatch][2][ibox].item<int>(),
 h = scores[ibatch][3][ibox].item<int>();
 unsigned short
 x1 = cx - w / 2,
 y1 = cy - h / 2,
 x2 = cx + w / 2,
 y2 = cy + h / 2;
 
 segmentation.push_back({ static_cast<int>(conf*100), x1, y1, x2, y2 , k_max });
 }
 }
 }

 std::sort(segmentation.begin(), segmentation.end(), [](const std::vector<int>& b1, const std::vector<int>& b2) {return b1[0] > b2[0]; });
 unsigned short
 x1 = static_cast<unsigned short>(segmentation[0][1]),
 y1 = static_cast<unsigned short>(segmentation[0][2]),
 x2 = static_cast<unsigned short>(segmentation[0][3]),
 y2 = static_cast<unsigned short>(segmentation[0][4]);
 int k_max = static_cast<int>(segmentation[0][5]);

 for (int j = 0; j < 160; j++)
 for (int i = 0; i < 160; i++)
 {
 ((float*)img160.data)[i + j * 160] = descriptors[0][k_max][j][i].item<float>();
 }

 img160.convertTo(img160u, CV_8U, 255, 8);
 cv::Mat imgMask(640, 640, CV_8U);
 cv::resize(img160u, imgMask, imgMask.size(), 0, 0, cv::INTER_LINEAR);
 cv::Mat imgMaskInDetection(y2 - y1, x2 - x1, CV_8U);
 cv::Mat outputIMG(y2 - y1, x2 - x1, CV_8UC3);
 cv::Vec3b Background(255, 255, 255);

 for (int i = 0; i < (y2 - y1); i++) {
 for (int j = 0; j < (x2 - x1); j++) {
 imgMaskInDetection.at<unsigned char>(i, j) = imgMask.at<unsigned char>(y1 + i, x1 + j);
 }
 }
 
 for (int i = 0; i < (y2 - y1); i++) {
 for (int j = 0; j < (x2 - x1); j++) {
 outputIMG.at<cv::Vec3b>(i, j) = _inputImage.at<cv::Vec3b>(y1 + i, x1 + j);
 }
 }
 
 for (int i = 0; i < imgMaskInDetection.rows; i++) {
 for (int j = 0; j < imgMaskInDetection.cols; j++) {
 unsigned char Maskpixel = imgMaskInDetection.at<unsigned char>(i, j);
 if ((int)Maskpixel == 0) {
 outputIMG.at<cv::Vec3b>(i, j) = Background;
 }
 }
 }
 
 return outputIMG;

 }

 cv::Mat ModelTorch::PrepImg(cv::Mat& input_img, int inputX, int inputY)
 {

 int imgSizeX = inputX;
 int imgSizeY = inputY;
 int halfImgSizeX = imgSizeX / 2;
 int halfImgSizeY = imgSizeY / 2;

 cv::Mat white_rectangel(imgSizeX, imgSizeY, CV_8UC3, cv::Scalar(255, 255, 255));

 cv::Mat img;
 input_img.copyTo(img);

 if (input_img.empty()) {
 std::cerr << "Error: Could not load the image.\n";
 return input_img;
 }

 cv::Size img_size = img.size();
 int img_w = img_size.width;
 int img_h = img_size.height;

 if (img_w < imgSizeX && img_h < imgSizeY) {

 cv::Mat new_img;
 white_rectangel.copyTo(new_img);



 int x = halfImgSizeX - img_w / 2;
 int y = halfImgSizeY - img_h / 2;

 cv::Mat imageROI = new_img(cv::Rect(x, y, img.cols, img.rows));

 img.copyTo(imageROI);


 return new_img;
 }

 if (img_w >= imgSizeX && img_h >= imgSizeY)
 {
 cv::Mat new_img = img(cv::Rect(img_w / 2 - halfImgSizeX, img_h / 2 - halfImgSizeY, imgSizeX, imgSizeY));

 return new_img;
 }

 if (img_w < imgSizeX || img_h < imgSizeY)
 {
 cv::Mat new_img;
 white_rectangel.copyTo(new_img);



 if (img_w < imgSizeX)
 {
 img = img(cv::Rect(0, img_h / 2 - halfImgSizeY, img_w, imgSizeY));
 }
 else
 {
 img = img(cv::Rect(img_w / 2 - halfImgSizeX, 0, imgSizeX, img_h));
 }

 img_size = img.size();
 img_w = img_size.width;
 img_h = img_size.height;

 int x = halfImgSizeX - img_w / 2;
 int y = halfImgSizeY - img_h / 2;

 cv::Mat imageROI = new_img(cv::Rect(x, y, img.cols, img.rows));

 img.copyTo(imageROI);


 return new_img;

 }
 }

} //namespase Bublik {